#!/bin/sh

set -eu

if [ -f "/usr/libexec/os-helpers-logging" ]; then
	# shellcheck disable=SC1091
	. /usr/libexec/os-helpers-logging
fi

INFO() {
	if command -v info > /dev/null; then
		info "$@"
	else
		echo "$@"
	fi
}

ERROR() {
	if command -v error > /dev/null; then
		error "$@"
	else
		echo "$@"
	fi
}

run_current_hooks_and_recover () {
	if [ "$hooks_rollback" = 1 ]; then
		# Run the current ones to cleanup the system.
		ERROR "Failed to run the new hooks. Running current hooks... "
		DURING_UPDATE=0 hostapp-update-hooks || true
	else
		ERROR "Failed to run the new hooks... "
	fi
	# unmount if we have mounted the new rootfs
	umount $new_rootfs || true
	ERROR "Update failed."
	exit 1
}

local_image=""
remote_image=""
reboot=0
hooks=1
hooks_rollback=1

while getopts 'f:i:rnx' flag; do
	case "${flag}" in
	f) local_image=$(realpath "${OPTARG}") ;;
	i) remote_image="${OPTARG}" ;;
	r) reboot=1 ;;
	n) hooks=0 ;;
	x) hooks_rollback=0 ;;
	*) error "Unexpected option ${flag}" ;;
	esac
done

if [ "$local_image" =  "" ] && [ "$remote_image" = "" ]; then
	ERROR "At least one of -f or -i is required"
	exit 1
fi

INFO "Running hostapp update..."

export DOCKER_HOST="unix:///var/run/balena-host.sock"
SYSROOT="/mnt/sysroot/inactive"
LOADTMP="/mnt/data/resin-data/tmp"

if [ -d /mnt/state ]; then
	# Save VPN state for rollbacks
	if [ -e /var/run/openvpn/vpn_status/active ]; then
		echo "BALENAOS_ROLLBACK_VPNONLINE=1" > /mnt/state/rollback-health-variables ;
	else
		echo "BALENAOS_ROLLBACK_VPNONLINE=0" > /mnt/state/rollback-health-variables ;
	fi
	sync -f /mnt/state
fi

# Initialize sysroot
mkdir -p /mnt/sysroot
for dir in 'dev' 'etc' 'balena' 'hostapps' 'mnt/state' 'proc' 'run' 'sbin' 'sys' 'tmp'; do
	mkdir -p "$SYSROOT/$dir"
done
if [ ! -f "$SYSROOT/etc/machine-id" ]; then
	touch "$SYSROOT/etc/machine-id"
fi
if [ ! -L "$SYSROOT/sbin/init" ]; then
	ln -sf ../current/boot/init "$SYSROOT/sbin/init"
fi
if [ ! -L "$SYSROOT/boot" ]; then
	ln -sf current/boot "$SYSROOT/boot"
fi

# Remove previous hostapp
for container in $(balena ps --all --quiet); do
	_out="$(balena rm --force --volumes "$container" 2>&1)"
	INFO "${_out}"
done
for image in $(balena images --all --quiet); do
	_out="$(balena rmi --force "$image" 2>&1)"
	INFO "${_out}"
done
for hostapp in "$SYSROOT/hostapps/"*; do
	[ -e "$hostapp" ] || break
	rm -rf "$hostapp"
done

# Load new hostapp
if [ "$local_image" != "" ]; then
	# bind mount the data partition for temporary extract/load files
	if ! mountpoint "${SYSROOT}/balena/tmp" >/dev/null; then
		mkdir -p "${LOADTMP}"
		mount --bind "${LOADTMP}" "${SYSROOT}/balena/tmp"
	fi

	# load a local image tarball
	HOSTAPP_IMAGE=$(balena load --quiet -i "$local_image" | cut -d: -f1 --complement | tr -d ' ')

	# attempt to unmount and clean up the temporary extract/load files but ignore failures
	umount "${SYSROOT}/balena/tmp" || true
	rm -rf "${LOADTMP}" || true
elif [ "$remote_image" != "" ]; then
	HOSTAPP_IMAGE="$remote_image"
	_out="$(balena pull "$HOSTAPP_IMAGE" 2>&1)"
	INFO "${_out}"
fi
CONTAINER_ID=$(balena create --runtime="bare" --volume=/boot "$HOSTAPP_IMAGE" /bin/sh)
BOOTSTRAP=$(balena inspect -f "{{range .Mounts}}{{.Destination}} {{.Source}}{{end}}" "$CONTAINER_ID" | awk '$1 == "/boot" { print $2 }' | head -n1)

# Create boot entry
rm -rf "$SYSROOT/hostapps/.new"
mkdir -p "$SYSROOT/hostapps/.new"
ln -sr "$BOOTSTRAP" "$SYSROOT/hostapps/.new/boot"
sync -f "$SYSROOT"
mv -T "$SYSROOT/hostapps/.new" "$SYSROOT/hostapps/$CONTAINER_ID"
sync -f "$SYSROOT"

# Mark it as current hostapp
ln -srf "$SYSROOT/hostapps/$CONTAINER_ID" "$SYSROOT/current.new"
sync -f "$SYSROOT"
mv -T "$SYSROOT/current.new" "$SYSROOT/current"
sync -f "$SYSROOT"

if [ "$hooks" = 1 ]; then
	# Run before from the new OS we are updating to in existing OS environment
	new_rootfs=$($SYSROOT/current/boot/init -sysroot $SYSROOT)
	if [ -f "$new_rootfs/usr/bin/hostapp-update-hooks-v2" ]; then
		if DURING_UPDATE=1 $new_rootfs/usr/bin/hostapp-update-hooks-v2 --before --dir "$new_rootfs/etc/hostapp-update-hooks.d/" ; then
			INFO "Before hooks (old os) ran successfully"
		else
			run_current_hooks_and_recover
		fi
	fi
	umount $new_rootfs

	# Run the defined hooks in the host OS we update to in new OS environment
	if _out=$(balena run --privileged --rm -v /dev:/dev -v /mnt:/mnt -v /sys:/sys -e DURING_UPDATE=1 "$HOSTAPP_IMAGE" hostapp-update-hooks 2>&1); then
		INFO "${_out}"
		INFO "New/Forward hooks (new os container) ran successfully."
	else
		ERROR "Failed running new/forward hooks in new os container: ${_out}"
		run_current_hooks_and_recover
	fi

	# Run after from the new OS we are updating to in existing OS environment
	new_rootfs=$($SYSROOT/current/boot/init -sysroot $SYSROOT)
	if [ -f "$new_rootfs/usr/bin/hostapp-update-hooks-v2" ]; then
		if DURING_UPDATE=1 $new_rootfs/usr/bin/hostapp-update-hooks-v2 --after --dir "$new_rootfs/etc/hostapp-update-hooks.d/"; then
			INFO "After hooks (old os) ran successfully"
		else
			run_current_hooks_and_recover
		fi
	fi
	umount $new_rootfs
fi

# Mark it as current partition
cur_counter=0
if [ -f "/mnt/sysroot/active/counter" ]; then
	cur_counter=$(cat /mnt/sysroot/active/counter)
fi
echo $((cur_counter + 1)) > "$SYSROOT/counter.new"
sync -f "$SYSROOT"
mv "$SYSROOT/counter.new" "$SYSROOT/counter"
sync -f "$SYSROOT"

INFO "Finished running hostapp update"

if [ "$reboot" = 1 ]; then
	if [ -x "/usr/libexec/safe_reboot" ]; then
		/usr/libexec/safe_reboot
	else
		reboot
	fi
fi
